/* *****************************************************************************
 * JFab (http://code.google.com/p/jfab)
 * Copyright (c) 2011 JFab.org
 * Admin jfab@jeffreythompson.net
 *
 * See the file "LICENSE.txt" for information on usage and redistribution of
 * this file, and for a DISCLAIMER OF ALL WARRANTIES.
 * *****************************************************************************/


package org.jfab.core.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import java.net.URI;
import java.net.URISyntaxException;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import javax.tools.DiagnosticCollector;
import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaCompiler;
import javax.tools.JavaCompiler.CompilationTask;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.JavaFileObject.Kind;
import javax.tools.SimpleJavaFileObject;
import javax.tools.StandardLocation;
import javax.tools.ToolProvider;


/**
 * Compile a String or other {@link CharSequence}, returning a Java {@link
 * Class} instance that may be instantiated. This class is a Facade around
 * {@link JavaStringCompiler} for a narrower use case, but a bit easier to use.
 *
 * <p>To compile a String containing source for a Java class which implements
 * MyInterface:
 *
 * <pre>
   ClassLoader classLoader = MyClass.class.getClassLoader(); // optional; null is also OK
   JavaStringCompiler&lt;Object&gt; compiler = new JavaStringCompiler&lt;MyInterface&gt;(classLoader, null);
   DiagnosticCollector&lt;JavaFileObject&gt; diagnostics = new DiagnosticCollector&lt;JavaFileObject&gt;(); // optional; null is also OK
   try
   {
      Class&lt;MyInterface&gt; newClass = compiler.compile(&quot;com.mypackage.NewClass&quot;,
            stringContaininSourceForNewClass, diagnostics, MyInterface);
      MyInterface instance = newClass.newInstance();
      instance.someOperation(someArgs);
   }
   catch (JavaStringCompilerException e)
   {
      handle(e);
   }
   catch (IllegalAccessException e)
   {
      handle(e);
   }
 * </pre>
 *
 * The source can be in a String, {@link StringBuffer}, or your own class which
 * implements {@link CharSequence}. If you implement your own, it must be thread
 * safe (preferably, immutable.)
 *
 * @param  <T>  Type.
 *
 * @see    <a
 *         href="http://www.ibm.com/developerworks/java/library/j-jcomp/index.html">
 *         Create dynamic applications with javax.tools</a>
 */
public final class JavaStringCompiler<T> implements Compiler<T>
{
    /** Compiler requires source files with a ".java" extension: */
    private static final String JAVA_EXTENSION = ".java";

    /** Class loader. */
    private final ClassLoaderImpl _classLoader;

    /** The compiler instance that this facade uses. */
    private final JavaCompiler _compiler;

    /** Collect compiler diagnostics in this instance. */
    private DiagnosticCollector<JavaFileObject> _diagnostics;

    /** The FileManager which will store source and class "files". */
    private final FileManagerImpl _javaFileManager;

    /** The compiler options (such as "-target" "1.5"). */
    private final List<String> _options;

    /**
     * Construct a new instance which delegates to the named class loader.
     *
     * @param   loader   the application ClassLoader. The compiler will look
     *                   through to this // class loader for dependent classes
     * @param   options  The compiler options (such as "-target" "1.5"). See the
     *                   usage for javac
     *
     * @throws  IllegalStateException  if the Java compiler cannot be loaded.
     */
    public JavaStringCompiler(ClassLoader loader, Iterable<String> options)
    {
        _compiler = ToolProvider.getSystemJavaCompiler();

        if (_compiler == null)
        {
            throw new IllegalStateException(
                "Cannot find the system Java compiler. "
                + "Check that your class path includes tools.jar");
        }

        _classLoader = new ClassLoaderImpl(loader);
        _diagnostics = new DiagnosticCollector<JavaFileObject>();

        final JavaFileManager fileManager = _compiler.getStandardFileManager(
                _diagnostics, null, null);

        // Create our FileManager which chains to the default file manager and our ClassLoader.
        _javaFileManager = new FileManagerImpl(fileManager, _classLoader);
        _options = new ArrayList<String>();

        if (options != null)
        {
            // Make a save copy of input options.
            for (String option : options)
            {
                _options.add(option);
            }
        }
    }

    /**
     * Factory method.
     *
     * @return  a new Java string compiler.
     */
    public final static Compiler<?> create()
    {
        ClassLoader classLoader = JavaStringCompiler.class.getClassLoader();
        List<String> options = Arrays.asList(new String[] { "-target", "1.6" });

        return new JavaStringCompiler<Object>(classLoader, options);
    }

    @Override public final synchronized Map<String, Class<T>> compile(
        final Map<String, CharSequence> classes,
        final DiagnosticCollector<JavaFileObject> diagnostics)
        throws JavaStringCompilerException
    {
        List<JavaFileObject> sources = new ArrayList<JavaFileObject>();

        for (Entry<String, CharSequence> entry : classes.entrySet())
        {
            String qualifiedClassName = entry.getKey();
            CharSequence javaSource = entry.getValue();

            if (javaSource != null)
            {
                final int dotPos = qualifiedClassName.lastIndexOf('.');
                final String className = (dotPos == -1)
                    ? qualifiedClassName
                    : qualifiedClassName.substring(dotPos + 1);
                final String packageName = (dotPos == -1)
                    ? "" : qualifiedClassName.substring(0, dotPos);
                final JavaFileObjectImpl source = new JavaFileObjectImpl(
                        className, javaSource);
                sources.add(source);
                // Store the source file in the FileManager via package/class
                // name.
                // For source files, we add a .java extension
                _javaFileManager.putFileForInput(StandardLocation.SOURCE_PATH,
                    packageName, className + JAVA_EXTENSION, source);
            }
        }

        // Get a CompliationTask from the compiler and compile the sources
        final CompilationTask task = _compiler.getTask(null, _javaFileManager,
                _diagnostics, _options, null, sources);
        final Boolean result = task.call();

        if ((result == null) || !result.booleanValue())
        {
            throw new JavaStringCompilerException("Compilation failed.",
                classes.keySet(), _diagnostics);
        }

        try
        {
            // For each class name in the inpput map, get its compiled
            // class and put it in the output map
            Map<String, Class<T>> compiled = new HashMap<String, Class<T>>();

            for (String qualifiedClassName : classes.keySet())
            {
                final Class<T> newClass = loadClass(qualifiedClassName);
                compiled.put(qualifiedClassName, newClass);
            }

            return compiled;
        }
        catch (ClassNotFoundException e)
        {
            throw new JavaStringCompilerException(classes.keySet(), e,
                _diagnostics);
        }
        catch (IllegalArgumentException e)
        {
            throw new JavaStringCompilerException(classes.keySet(), e,
                _diagnostics);
        }
        catch (SecurityException e)
        {
            throw new JavaStringCompilerException(classes.keySet(), e,
                _diagnostics);
        }
    }

    @Override public final synchronized Class<T> compile(
        final String qualifiedClassName, final CharSequence javaSource,
        final DiagnosticCollector<JavaFileObject> diagnostics,
        final Class<?>... types) throws JavaStringCompilerException,
        ClassCastException
    {
        if (diagnostics != null)
        {
            _diagnostics = diagnostics;
        }
        else
        {
            _diagnostics = new DiagnosticCollector<JavaFileObject>();
        }

        Map<String, CharSequence> classes = new HashMap<String, CharSequence>(
                1);
        classes.put(qualifiedClassName, javaSource);

        Map<String, Class<T>> compiled = compile(classes, diagnostics);
        Class<T> newClass = compiled.get(qualifiedClassName);
        return castable(newClass, types);
    }

    @Override public final ClassLoader getClassLoader()
    {
        return _javaFileManager.getClassLoader();
    }

    @Override
    @SuppressWarnings("unchecked")
    public final Class<T> loadClass(final String qualifiedClassName)
        throws ClassNotFoundException
    {
        return (Class<T>)_classLoader.loadClass(qualifiedClassName);
    }

    /**
     * Converts a String to a URI.
     *
     * @param   name  a file name
     *
     * @return  a URI
     */
    static final URI toURI(String name)
    {
        try
        {
            return new URI(name);
        }
        catch (URISyntaxException e)
        {
            throw new RuntimeException(e);
        }
    }

    /**
     * Check that the <var>newClass</var> is a subtype of all the type
     * parameters and throw a ClassCastException if not.
     *
     * @param   newClass  New class.
     * @param   types     zero of more classes or interfaces that the <var>
     *                    newClass</var> must be castable to.
     *
     * @return  <var>newClass</var> if it is castable to all the types
     *
     * @throws  ClassCastException  if <var>newClass</var> is not castable to
     *                              all the types.
     */
    private Class<T> castable(Class<T> newClass, Class<?>... types)
        throws ClassCastException
    {
        for (Class<?> type : types)
        {
            if (!type.isAssignableFrom(newClass))
            {
                throw new ClassCastException(type.getName());
            }
        }

        return newClass;
    }

    /**
     * A custom ClassLoader which maps class names to JavaFileObjectImpl
     * instances.
     */
    private static final class ClassLoaderImpl extends ClassLoader
    {
        /** Classes. */
        private final Map<String, JavaFileObject> _classes =
            new HashMap<String, JavaFileObject>();

        /**
         * Construct this object.
         *
         * @param  parentClassLoader  Parent class loader.
         */
        ClassLoaderImpl(final ClassLoader parentClassLoader)
        {
            super(parentClassLoader);
        }

        /**
         * @see  java.lang.ClassLoader#getResourceAsStream(java.lang.String)
         */
        @Override public InputStream getResourceAsStream(final String name)
        {
            if (name.endsWith(".class"))
            {
                String qualifiedClassName = name.substring(0,
                        name.length() - ".class".length()).replace('/', '.');
                JavaFileObjectImpl file = (JavaFileObjectImpl)_classes.get(
                        qualifiedClassName);

                if (file != null)
                {
                    return new ByteArrayInputStream(file.getByteCode());
                }
            }

            return super.getResourceAsStream(name);
        }

        /**
         * @see  java.lang.ClassLoader#findClass(java.lang.String)
         */
        @Override protected Class<?> findClass(final String qualifiedClassName)
            throws ClassNotFoundException
        {
            JavaFileObject file = _classes.get(qualifiedClassName);

            if (file != null)
            {
                byte[] bytes = ((JavaFileObjectImpl)file).getByteCode();
                return defineClass(qualifiedClassName, bytes, 0, bytes.length);
            }

            // Workaround for "feature" in Java 6
            // see http://bugs.sun.com/bugdatabase/view_bug.do?bug_id=6434149
            try
            {
                Class<?> c = Class.forName(qualifiedClassName);
                return c;
            }
            catch (ClassNotFoundException nf)
            {
                // Ignore and fall through
            }

            return super.findClass(qualifiedClassName);
        }

        /**
         * @see  java.lang.ClassLoader#loadClass(java.lang.String, boolean)
         */
        @Override protected synchronized Class<?> loadClass(final String name,
            final boolean resolve) throws ClassNotFoundException
        {
            return super.loadClass(name, resolve);
        }

        /**
         * Add a class name/JavaFileObject mapping
         *
         * @param  qualifiedClassName  the name
         * @param  javaFile            the file associated with the name
         */
        void add(final String qualifiedClassName, final JavaFileObject javaFile)
        {
            _classes.put(qualifiedClassName, javaFile);
        }

        /**
         * @return  An collection of JavaFileObject instances for the classes in
         *          the class loader.
         */
        Collection<JavaFileObject> files()
        {
            return Collections.unmodifiableCollection(_classes.values());
        }
    }

    /**
     * A JavaFileManager which manages Java source and classes. This FileManager
     * delegates to the JavaFileManager and the ClassLoaderImpl provided in the
     * constructor. The sources are all in memory CharSequence instances and the
     * classes are all in memory byte arrays.
     */
    private static final class FileManagerImpl
        extends ForwardingJavaFileManager<JavaFileManager>
    {
        /** The delegating class loader (passed to the constructor). */
        private final ClassLoaderImpl _classLoader;

        /** Internal map of filename URIs to JavaFileObjects. */
        private final Map<URI, JavaFileObject> _fileObjects =
            new HashMap<URI, JavaFileObject>();

        /**
         * Construct a new FileManager which forwards to the <var>
         * fileManager</var> for source and to the <var>classLoader</var> for
         * classes
         *
         * @param  fileManagerIn  another FileManager that this instance
         *                        delegates to for additional source.
         * @param  classLoader    a ClassLoader which contains dependent classes
         *                        that the compiled classes will require when
         *                        compiling them.
         */
        public FileManagerImpl(JavaFileManager fileManagerIn,
            ClassLoaderImpl classLoader)
        {
            super(fileManagerIn);
            _classLoader = classLoader;
        }

        /**
         * @return  the class loader which this file manager delegates to
         */
        public ClassLoader getClassLoader()
        {
            return _classLoader;
        }

        @Override public ClassLoader getClassLoader(
            JavaFileManager.Location location)
        {
            return _classLoader;
        }

        /**
         * For a given file <var>location</var>, return a FileObject from which
         * the compiler can obtain source or byte code.
         *
         * @param   location      an abstract file location
         * @param   packageName   the package name for the file
         * @param   relativeName  the file's relative name
         *
         * @return  a FileObject from this or the delegated FileManager
         *
         * @see     javax.tools.ForwardingJavaFileManager#getFileForInput(javax.tools.JavaFileManager.Location,
         *          java.lang.String, java.lang.String)
         */
        @Override public FileObject getFileForInput(Location location,
            String packageName, String relativeName) throws IOException
        {
            FileObject o = _fileObjects.get(uri(location, packageName,
                        relativeName));

            if (o != null)
            {
                return o;
            }

            return super.getFileForInput(location, packageName, relativeName);
        }

        /**
         * Create a JavaFileImpl for an output class file and store it in the
         * classloader.
         *
         * @see  javax.tools.ForwardingJavaFileManager#getJavaFileForOutput(javax.tools.JavaFileManager.Location,
         *       java.lang.String, javax.tools.JavaFileObject.Kind,
         *       javax.tools.FileObject)
         */
        @Override public JavaFileObject getJavaFileForOutput(Location location,
            String qualifiedName, Kind kind, FileObject outputFile)
            throws IOException
        {
            JavaFileObject file = new JavaFileObjectImpl(qualifiedName, kind);
            _classLoader.add(qualifiedName, file);
            return file;
        }

        /**
         * @see  javax.tools.ForwardingJavaFileManager#inferBinaryName(javax.tools.JavaFileManager.Location,
         *       javax.tools.JavaFileObject)
         */
        @Override public String inferBinaryName(Location loc,
            JavaFileObject file)
        {
            String result;

            // For our JavaFileImpl instances, return the file's name, else
            // simply run the default implementation
            if (file instanceof JavaFileObjectImpl)
            {
                result = file.getName();
            }
            else
            {
                result = super.inferBinaryName(loc, file);
            }

            return result;
        }

        /**
         * @see  javax.tools.ForwardingJavaFileManager#list(javax.tools.JavaFileManager.Location,
         *       java.lang.String, java.util.Set, boolean)
         */
        @Override public Iterable<JavaFileObject> list(Location location,
            String packageName, Set<Kind> kinds, boolean recurse)
            throws IOException
        {
            Iterable<JavaFileObject> result = super.list(location, packageName,
                    kinds, recurse);
            ArrayList<JavaFileObject> files = new ArrayList<JavaFileObject>();

            if ((location == StandardLocation.CLASS_PATH)
                && kinds.contains(JavaFileObject.Kind.CLASS))
            {
                for (JavaFileObject file : _fileObjects.values())
                {
                    if ((file.getKind() == Kind.CLASS)
                        && file.getName().startsWith(packageName))
                    {
                        files.add(file);
                    }
                }

                files.addAll(_classLoader.files());
            }
            else if ((location == StandardLocation.SOURCE_PATH)
                && kinds.contains(JavaFileObject.Kind.SOURCE))
            {
                for (JavaFileObject file : _fileObjects.values())
                {
                    if ((file.getKind() == Kind.SOURCE)
                        && file.getName().startsWith(packageName))
                    {
                        files.add(file);
                    }
                }
            }

            for (JavaFileObject file : result)
            {
                files.add(file);
            }

            return files;
        }

        /**
         * Store a file that may be retrieved later with {@link
         * #getFileForInput(javax.tools.JavaFileManager.Location, String,
         * String)}
         *
         * @param  location      the file location
         * @param  packageName   the Java class' package name
         * @param  relativeName  the relative name
         * @param  file          the file object to store for later retrieval
         */
        public void putFileForInput(StandardLocation location,
            String packageName, String relativeName, JavaFileObject file)
        {
            _fileObjects.put(uri(location, packageName, relativeName), file);
        }

        /**
         * Convert a location and class name to a URI
         *
         * @param   location      Location.
         * @param   packageName   Package name.
         * @param   relativeName  Relative name.
         *
         * @return  a new URI.
         */
        private URI uri(Location location, String packageName,
            String relativeName)
        {
            return JavaStringCompiler.toURI(location.getName() + '/'
                    + packageName + '/' + relativeName);
        }
    }

    /**
     * A JavaFileObject which contains either the source text or the compiler
     * generated class. This class is used in two cases.
     *
     * <ol>
     * <li>This instance uses it to store the source which is passed to the
     * compiler. This uses the {@link
     * JavaStringCompiler.JavaFileObjectImpl#JavaFileObjectImpl(String,
     * CharSequence)} constructor.
     * <li>The Java compiler also creates instances (indirectly through the
     * FileManagerImplFileManager) when it wants to create a JavaFileObject for
     * the .class output. This uses the {@link
     * JavaStringCompiler.JavaFileObjectImpl#JavaFileObjectImpl(String,
     * JavaFileObject.Kind)} constructor.
     * </ol>
     *
     * This class does not attempt to reuse instances (there does not seem to be
     * a need, as it would require adding a Map for the purpose, and this would
     * also prevent garbage collection of class byte code.)
     */
    private static final class JavaFileObjectImpl extends SimpleJavaFileObject
    {
        /** If kind == CLASS, this stores byte code from openOutputStream. */
        private ByteArrayOutputStream _byteCode;

        /** If kind == SOURCE, this contains the source text */
        private final CharSequence _source;

        /**
         * Construct a new instance
         *
         * @param  name    the file name
         * @param  kindIn  the kind of file
         */
        JavaFileObjectImpl(final String name, final Kind kindIn)
        {
            super(JavaStringCompiler.toURI(name), kindIn);
            _source = null;
        }

        /**
         * Construct a new instance which stores source
         *
         * @param  baseName  the base name
         * @param  source    the source code
         */
        JavaFileObjectImpl(final String baseName, final CharSequence source)
        {
            super(JavaStringCompiler.toURI(
                    baseName + JavaStringCompiler.JAVA_EXTENSION), Kind.SOURCE);
            _source = source;
        }

        /**
         * Return the source code content
         *
         * @see  javax.tools.SimpleJavaFileObject#getCharContent(boolean)
         */
        @Override public CharSequence getCharContent(
            final boolean ignoreEncodingErrors)
            throws UnsupportedOperationException
        {
            if (_source == null)
            {
                throw new UnsupportedOperationException("getCharContent()");
            }

            return _source;
        }

        /**
         * Return an input stream for reading the byte code
         *
         * @see  javax.tools.SimpleJavaFileObject#openInputStream()
         */
        @Override public InputStream openInputStream()
        {
            return new ByteArrayInputStream(getByteCode());
        }

        /**
         * Return an output stream for writing the bytecode
         *
         * @see  javax.tools.SimpleJavaFileObject#openOutputStream()
         */
        @Override public OutputStream openOutputStream()
        {
            _byteCode = new ByteArrayOutputStream();
            return _byteCode;
        }

        /**
         * @return  the byte code generated by the compiler
         */
        byte[] getByteCode()
        {
            return _byteCode.toByteArray();
        }
    }
}
